
import os
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from torchvision import transforms
from typing import Any, Callable, Dict, List, Optional, Tuple

# import torch.nn as nn
# import torch.nn.functional as F
# Fu.relu()是函数调用，一般使用在foreward函数里。
# nn.ReLU()是模块调用，一般在定义网络层的时候使用。



ROOT = "e:/sfxData/DeepLearning"

MNIST_TRANSFORM = transforms.Compose([
    transforms.ToTensor(),                      # 转换为张量
    # transforms.Normalize((0.1307,), (0.3081,))  # 设定标准化值
])


def get_dataset(root: str, train: bool = True, transform: Optional[Callable] = None) -> MNIST:
    """
    取得 MNIST 数据集
    """
    ds = _mnist(root, train=train, transform=transform)
    return ds


def _mnist(root: str, train: bool = True, transform: Optional[Callable] = None) -> MNIST:
    """
    取得 MNIST 数据集.

    Args
    :root - 数据集根目录
    :train - 是否取得训练集
    """
    xtransform = MNIST_TRANSFORM if transform == None else transform
    isDownload = not(os.path.exists(root)) or not os.listdir(root)
    dset = MNIST(root=root, train=train,
                 transform=xtransform, download=isDownload)
    return dset
